import logging

import numpy as np
from scipy import signal
from scipy import fft
from scipy import interpolate
from scipy.io import wavfile
import matplotlib.pyplot as plt

# Import everything from `lib/dsp.py`.
from lib.dsp import *


def upsample_filter(x, b, a, up_factor):
    """
    Parameters
    ---
    x: array_like
        低采样率信号
    b: array_like
        滤波器参数b
    a: array_like
        滤波器参数a
    up_factor: int
        增采样倍数
    Returns
    ---
    x_filter: array_like
        增采样后信号
    """
    # 在采样信号中插0
    # 注意: 最后一个点不需要插0
    x_insert_0 = np.zeros((len(x) - 1) * up_factor + 1)
    x_insert_0[::up_factor] = x

    # 获取滤波后序列
    x_filter = signal.lfilter(b, a, x_insert_0) * up_factor

    return x_filter


def evaluate_visual(x_upsample, x_target, fs_up):
    """
    Parameters
    ---
    x_upsample: array_like
        增采样信号有效区段
    x_target: array_like
        目标增采样信号有效区段
    fs_up: float
        增采样率后采样频率
    Returns
    ---
    x_upsample: array_like
        增采样信号有效区段时域序列
    x_upsample_fft: array_like
        增采样信号有效区段频谱序列
    x_target: array_like
        目标增采样信号有效区段时域序列
    x_target_fft: array_like
        目标增采样信号有效区段频谱序列
    """
    # Plot the signals for inspection.
    fig, axs = plt.subplots(1, 2, figsize=(12, 2), squeeze=False)
    TITLES = ['sequence', 'spectrum']
    plot_signals(
        ['x_target', 'x_upsample'],
        [x_target, x_upsample],
        fs_up, axs[0, :],
        TITLES
    )
    return x_upsample, np.fft.fft(x_upsample), x_target, np.fft.fft(x_target)


def evaluate_objective(x_upsample, x_target):
    """
    Parameters
    ---
    x_upsample: array_like
        增采样信号有效区段
    x_target: array_like
        目标增采样信号有效区段
    Returns
    ---
    snr: float
        信噪比
    """
    # 扩大数据类型长度，防止溢出
    x_upsample = x_upsample.astype(np.float64)
    x_target = x_target.astype(np.float64)
    # 计算信噪比
    return 10 * np.log10(np.sum(x_target**2) / np.sum((x_target - x_upsample)**2))


def build_filter(fs, kind='fir', R_p_target=1, A_s_target=40):
    """
    Parameters
    ---
    fs:int
        增采样后采样频率
    kind:str
        滤波器类型{fir, iir}
    R_p_target: float
        通带最大纹波
    A_s_target:
        阻带最小衰减
    Returns
    ---
    b:array_like
        滤波器b参数
    a:array_like
        滤波器a参数
    bands:array_like
        滤波器带信息
    """
    # 语音信号频率范围 300-3.4kHz
    w_p, w_st = 3.4e3, 4.6e3

    # 过渡带宽度
    f_tran = abs(w_p - w_st)
    # 过渡带宽度需表示为 过渡带带宽 / 奈奎斯特频率
    width = f_tran / (fs / 2)

    a = None
    b = None
    if kind == 'fir':
        # 使用kaiserord计算凯泽窗参数，第一个参数为阻带最小衰减(单位dB), 第二个参数为过渡带宽度
        N, beta = signal.kaiserord(A_s_target, width)
        window = ('kaiser', beta)

        # Compute the FIR coefficients.
        w_c = (w_p + w_st) / 2
        b = signal.firwin(N, w_c, window=window, fs=fs)
        a = [1]
    else:
        # ftype='butter'
        # ftype='cheby1'
        # ftype='cheby2'
        ftype = 'ellip'
        b, a = signal.iirdesign(
            w_p, w_st,
            R_p_target, A_s_target,
            ftype=ftype,
            fs=fs)

    # 滤波器带信息
    bands = [
        ('pass', 0, w_p),
        ('tran', w_p, w_st),
        ('stop', w_st, fs/2)
    ]

    return b, a, bands


def upsample_interp(x, up_factor, kind='nearest'):
    """
    Parameters
    ---
    x: array_like
        低采样率信号
    up_factor: int
        增采样倍数
    kind: str
        插值算法
    Returns
    ---
    x_interp: array_like
        增采样后信号
    """
    # 低采样率信号横坐标
    w = np.linspace(0, len(x), num=len(x), endpoint=False)
    # 生成插值函数
    f = interpolate.interp1d(w, x, kind=kind)

    # 增采样信号横坐标
    x_interp = np.linspace(0, len(x)-1, num=(len(x) - 1) * up_factor + 1)
    # 增采样信号
    x_interp = f(x_interp)
    return x_interp


def do_upsample(args, up_factor=6):
    # 输出命令行参数
    logging.info(
        f'SCHEMA:{args.SCHEMA}, INPUT:{args.INPUT}, OUTPUT:{args.OUTPUT}')

    # 读入WAV文件信息
    rate, data = wavfile.read(args.INPUT)
    # 输出WAV文件采样率
    logging.info(f'Input wav file\'s sample rate is {rate}')

    # 增采样后采样频率
    fs_up = rate * up_factor

    upsampled = None
    # 选择增采样方法
    if args.SCHEMA == 'nearest' or args.SCHEMA == 'linear':
        # 时域增采样
        upsampled = upsample_interp(data, up_factor, args.SCHEMA)
    else:
        # 滤波器类型
        filter_type = args.SCHEMA[:3]
        # 阻带最小衰减
        A_s_target = int(args.SCHEMA[4:])
        logging.info(f'Filter type:{filter_type}, A_s_target:{A_s_target}')
        # 滤波增采样
        b, a, bands = build_filter(
            fs_up, kind=filter_type, A_s_target=A_s_target)
        upsampled = upsample_filter(data, b, a, up_factor)

    wavfile.write(args.OUTPUT, fs_up, upsampled.astype(data.dtype))
    logging.info(
        f'Saved upsampled wav file to: {args.OUTPUT}, Upsampled Rate:{fs_up}')


def append_to_csv_by_row(path, row, mode='a+'):
    """
    将指定行附加到csv文件末尾
    :param path: csv文件路径
    :param row: 要附加的行
    :param mode: 文件打开模式
    :return: None
    """
    import csv
    # 以附加模式打开文件
    with open(path, mode, newline='') as csv_file:
        csv_writer = csv.writer(csv_file, quoting=csv.QUOTE_ALL)
        # 写入一行
        csv_writer.writerow(row)


def do_snr(args):
    logging.info(
        f'UPSAMPLED:{args.UPSAMPLED}, TARGET:{args.TARGET}, START1:{args.START1}, START2:{args.START2}, LENGTH:{args.LENGTH}')
    # 读取增采样后数据
    rate_up, data_up = wavfile.read(args.UPSAMPLED)
    # 读取目标增采样数据
    rate_target, data_targt = wavfile.read(args.TARGET)
    # 输出文件采样率
    logging.info(f'Upsampled Rate:{rate_up}, Target Rate:{rate_target}')

    # 计算信噪比
    snr = evaluate_objective(data_up[args.START1:args.START1+args.LENGTH],
                             data_targt[args.START2:args.START2+args.LENGTH])
    logging.info(f'SNR: {snr} dB')

    if args.csv:
        # 写入到CSV文件
        append_to_csv_by_row(args.csv, [
                             args.UPSAMPLED,
                             args.TARGET,
                             args.START1,
                             args.START2,
                             args.LENGTH,
                             snr])


def parse_args(args=None):
    """
    处理命令行输入

    参数:
        audioUpsample upsample SCHEMA INPUT OUTPUT
        audioUpsample snr UPSAMPLED TARGET START1 START2 LENGTH
    """
    import argparse
    import pathlib

    # 编解码命令解析
    operation_parser = argparse.ArgumentParser()
    # test 参数, 是否运行自动测试
    operation_parser.add_argument('-t', '--test', action="store_true",
                                  help='run tests')
    # verbose 参数, 用于控制log等级
    operation_parser.add_argument('-v', '--verbose', action="store_true",
                                  help='show debug message')
    # 置入空到func中, 否则在没有输入参数的情况下会报变量不存在的错
    operation_parser.set_defaults(func=None)

    # 添加子命令
    subparsers = operation_parser.add_subparsers()

    # 增采样子命令
    upsample_parser = subparsers.add_parser(
        'upsample', help='upsample input wav file with specificed schema')
    # 增采样方案参数
    upsample_parser.add_argument(
        'SCHEMA', help='schema for upsampling the input wav', type=str)
    # 输入WAV文件
    upsample_parser.add_argument(
        'INPUT', help='low sample rate wav file', type=pathlib.Path)
    # 输出WAV文件
    upsample_parser.add_argument(
        'OUTPUT', help='upsampled wav file', type=pathlib.Path)
    # 设置处理方法
    upsample_parser.set_defaults(func=do_upsample)

    # SNR计算子命令
    snr_parser = subparsers.add_parser(
        'snr', help='calculate the SNR of the wav signal')
    # 增采样后文件
    snr_parser.add_argument(
        'UPSAMPLED', help='Upsampled wav file', type=pathlib.Path)
    # 目标增采样文件
    snr_parser.add_argument(
        'TARGET', help='target upsampled wav file', type=pathlib.Path)
    # 增采样文件起始
    snr_parser.add_argument(
        'START1', help='start offset of the upsampled wav file', type=int)
    # 目标增采样文件起始
    snr_parser.add_argument(
        'START2', help='start offset of the target upsampled wav file', type=int)
    # 计算SNR的长度
    snr_parser.add_argument('LENGTH', help='length to calc the SNR', type=int)
    # 存储SNR数据的CSV文件
    snr_parser.add_argument(
        '--csv', help='path to csv file to save snr', type=pathlib.Path)
    # 设置处理方法
    snr_parser.set_defaults(func=do_snr)

    # 处理输入的命令
    args = operation_parser.parse_args(args)

    # 配置日志输出等级
    logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO)

    if args.func:
        # 运行子命令
        args.func(args)
    else:
        # 输出帮助
        operation_parser.print_help()


if __name__ == '__main__':
    import itertools
    for filter_type, A_s_target in itertools.product(['fir', 'iir'], [20, 40, 70]):
        schema = f'{filter_type}-{A_s_target}'
        out_name = f'data/p232_045.48KHz-{filter_type}-{A_s_target}.wav'
        parse_args(['upsample', schema, 'data/p232_045.8KHz.wav', out_name])

        # 根据参数构造滤波器
        b, a, bands = build_filter(48e3, filter_type, A_s_target=A_s_target)

        # 取出通带信息
        band_type, w_start, w_end = bands[0]
        assert band_type == 'pass', 'Only support low pass filter'
        # 计算通带群延迟
        _, gd_tmp = signal.group_delay(
            (b, a),
            w=np.linspace(w_start, w_end, 512),
            fs=48e3)
        # 将通带群延迟取平均
        gd = int(np.round(gd_tmp.mean()))

        # 计算SNR的长度
        _, data = wavfile.read(out_name)
        length = len(data) - gd

        parse_args(
            ['snr', '--csv', 'data/snr.csv', out_name, 'data/p232_045.48KHz.wav', f'{gd}', f'{0}', f'{length}'])
    for schema in ['nearest', 'linear']:
        out_name = f'data/p232_045.48KHz-{schema}.wav'
        parse_args(['upsample', schema, 'data/p232_045.8KHz.wav', out_name])

        # 计算SNR的长度
        _, data = wavfile.read(out_name)
        length = len(data)

        parse_args(
            ['snr', '--csv', 'data/snr.csv', out_name, 'data/p232_045.48KHz.wav', f'{0}', f'{0}', f'{length}'])
